SparkのUnitTest作成でspark-testing-baseを使うメリットとimport spark.implicit._について辿ってみた

Sparkを使ったテストを書く上で便利なspark-testing-baseについての挙動や、よく見かける「import spark.implicit._」が何をしているのかを追ってみました。

#Scala

#Apache Spark

haoyayoi

2020.06.25

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

sparkを使ったコードでテストを作成する際、assert用のDataFrameやRowの生成を行うロジックを自前で持つと検証コストが非常に高くなります。

シンプルにsparkのテストを書くことを目的とした spark-testing-base を用いてのテスト実装のメリットについて書いてみました。

build.sbtに追加する

記事執筆時点でのバージョンです。必要に応じたものを指定してください。

libraryDependencies ++= Seq(
  ..
  "com.holdenkarau" %% "spark-testing-base" % "2.4.3_0.14.0" % "test",
)

spark-testing-baseでのテスト

assert文のためにDataFrameやRowを生成することには変わりありませんが、sparkに用意されているspark.implicits._を利用する事で簡潔になります。

import com.holdenkarau.spark.testing.DataFrameSuiteBase
import org.scalatest.funsuite.AnyFunSuite
import org.scalatest.matchers.should.Matchers

class MyTest extends AnyFunSuite with Matchers with DataFrameSuiteBase {
  test("test-1") {
    import spark.implicits._
    val df = List(("1", "Tarou")).toDF("id", "name")
    df.collect().length should be (1)
    val row = df.first()
    row.get(0) should be ("1")
  }
}

SparkのSessionを管理するコードが見当たらず、かつListを用いてDataFrameを生成するという離れ技で、仕組みが分かっていないと何故これがグリーンになるのか理解できない状態になるかもしれません。

DataFrameSuiteBaseを利用してSession管理を自動化する

spark-testing-baseからDataFrameSuiteBaseをmixinすると自動でSessionが生成されるようになり、テストケース内でのSession管理を意識する必要がなくなります。

以下のように事前にSparkSessionを生成するコード例も見かけましたが、DataFrameSuiteBaseを利用することで不要になります。

val spark = val spark = SparkSession.builder().getOrCreate()
import spark.implicits._

また、DataFrameSuiteBaseは内部でScalaTestのbeforeAllやafterAllに合わせたContextの自動管理も行っており、ScalaTestを利用してのSpark向けテストを書くのであれば断然利用しない手はありません。

trait DataFrameSuiteBase extends TestSuite with SharedSparkContext with DataFrameSuiteBaseLike { self: Suite =>
  
  override def beforeAll() {
    super.beforeAll()
    super.sqlBeforeAllTestCases()
  }
  
  override def afterAll() {
    super.afterAll()
    SQLContextProvider._sqlContext = null
  }
}

spark.implicits._ をimportして暗黙のDataSet生成を行う

import spark.implicits._を追記するだけでListからDataFrameが生成出来るようになります。種はSparkSession.scala内の実装にあります。

  object implicits extends SQLImplicits with Serializable {
    protected override def _sqlContext: SQLContext = SparkSession.this.sqlContext
  }

implicitsが継承しているSQLImplicitsは基本的なScalaオブジェクトをDataSetにする暗黙の手続き集となっています。

A collection of implicit methods for converting common Scala objects into [[Dataset]]s.

なお、importの対象にspark.implicits._が含まれる事は認識し難いのですが、spark-shellを利用することで存在することがわかるようになっています。

% brew install apache-spark
% spark-shell

20/06/25 15:55:46 WARN Utils: Your hostname, localhost resolves to a loopback address: 127.0.0.1; using 192.168.0.186 instead (on interface en0)
20/06/25 15:55:46 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address
20/06/25 15:55:46 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
Setting default log level to "WARN".
To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).
Spark context Web UI available at http://192.168.0.186:4040
Spark context available as 'sc' (master = local[*], app id = local-1593068152354).
Spark session available as 'spark'.
Welcome to
      ____              __
     / __/__  ___ _____/ /__
    _\ \/ _ \/ _ `/ __/  '_/
   /___/ .__/\_,_/_/ /_/\_\   version 2.4.6
      /_/

Using Scala version 2.11.12 (OpenJDK 64-Bit Server VM, Java 1.8.0_252)
Type in expressions to have them evaluated.
Type :help for more information.

scala> :imports
 1) import org.apache.spark.SparkContext._ (70 terms, 1 are implicit)
 2) import spark.implicits._       (1 types, 67 terms, 37 are implicit)
 3) import spark.sql               (1 terms)
 4) import org.apache.spark.sql.functions._ (385 terms)